package com.dtdream.emr;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDTF;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.WritableConstantIntObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.WritableConstantStringObjectInspector;

public class TransArray extends GenericUDTF {
    private ArrayList<PrimitiveObjectInspector> stringOI = null;
    public static int keyNum = 1;
    public static String seperator = "";

    @Override
    public StructObjectInspector initialize(ObjectInspector[] args) throws UDFArgumentException {
        if (args.length <= 2) {
            throw new UDFArgumentException("TransArray() takes more than two argument");
        }

        if (args[0].getCategory() != ObjectInspector.Category.PRIMITIVE
                && ((PrimitiveObjectInspector) args[0]).getPrimitiveCategory() != PrimitiveObjectInspector.PrimitiveCategory.INT) {
            throw new UDFArgumentException("TransArray() takes a string as first parameter");
        }
        if (args[1].getCategory() != ObjectInspector.Category.PRIMITIVE
                && ((PrimitiveObjectInspector) args[1]).getPrimitiveCategory() != PrimitiveObjectInspector.PrimitiveCategory.STRING) {
            throw new UDFArgumentException("TransArray() takes a string as second parameter");
        }

        // input
        TransArray.keyNum = ((WritableConstantIntObjectInspector) args[0]).getWritableConstantValue().get();
        TransArray.seperator = ((WritableConstantStringObjectInspector) args[1]).getWritableConstantValue().toString();
        stringOI = new ArrayList<PrimitiveObjectInspector>();

        for (int i = 2; i < args.length; i++) {
            stringOI.add((PrimitiveObjectInspector) args[i]);
        }


        // output
        List<String> fieldNames = new ArrayList<String>(args.length - 2);
        List<ObjectInspector> fieldOIs = new ArrayList<ObjectInspector>(args.length - 2);

        for (int i = 2; i < args.length; i++) {
            fieldNames.add("c");
        }
        for (int i = 2; i < args.length; i++) {
            fieldOIs.add(PrimitiveObjectInspectorFactory.javaStringObjectInspector);
        }
        return ObjectInspectorFactory.getStandardStructObjectInspector(fieldNames, fieldOIs);
    }

    public ArrayList<Object[]> processInputRecord(ArrayList<String> oriRecord) {
        ArrayList<Object[]> result = new ArrayList<Object[]>();
        ArrayList<String[]> splitCoulmn = new ArrayList<>();
        int keyNum = TransArray.keyNum;
        String separator = TransArray.seperator;
        int maxlength = 0;

        for (int i = keyNum; i < oriRecord.size(); i++) {
            if (null == oriRecord.get(i)){
                continue;
            }

            String[] tmp = oriRecord.get(i).split(separator);
            if (tmp.length > maxlength) {
                maxlength = tmp.length;
            }
            splitCoulmn.add(tmp);
        }
        for (int i = 0; i < maxlength; i++) {
            Object[] record = new Object[oriRecord.size()];
            for (int j = 0; j < keyNum; j++) {
                record[j] = oriRecord.get(j);
            }
            for (int j = keyNum; j < oriRecord.size(); j++) {
                if (splitCoulmn.get(j - keyNum).length > i) {
                    record[j] = splitCoulmn.get(j - keyNum)[i];
                } else {
                    record[j] = "NULL";
                }
            }
            result.add(record);
        }
        return result;
    }

    @Override
    public void process(Object[] record) throws HiveException {
        if (null == record){
            return;
        }
        final ArrayList<String> oriRecord = new ArrayList<String>(record.length);
        for (int i = 2; i < record.length; i++) {
            if (null == stringOI || null == stringOI.get(i - 2) || null == record[i]){
                continue;
            }
            Object obj = stringOI.get(i - 2).getPrimitiveJavaObject(record[i]);
            if (null != obj){
                oriRecord.add(stringOI.get(i - 2).getPrimitiveJavaObject(record[i]).toString());
            }
        }

        ArrayList<Object[]> results = processInputRecord(oriRecord);
        Iterator<Object[]> it = results.iterator();
        while (it.hasNext()) {
            Object[] r = it.next();
            forward(r);
        }
    }

    @Override
    public void close() throws HiveException {
        // do nothing
    }
}